#!/opt/local/bin/python

from scipy import pi, log10, exp, arctan
from numpy import logspace, arange, rad2deg, unwrap
import matplotlib.pyplot as plt

def phase(x):
    ini = arctan(x.imag/x.real)
    if ini>0:
        if x.imag<0 and x.real<0:
            return ini - pi
        else:
            return ini
    elif ini<0:
        if x.real<0 and x.imag>0:
            return ini + pi
        else:
            return ini
    else: # ini==0 presumably
        return ini

class control(object):
    def __init__(self, Kp, Ki, tau):
        self.Kp  = Kp
        self.Ki  = Ki
        self.tau = tau

    def getParams(self):
        return (self.Kp, self.Ki, self.tau)

class cavBox(object):
    def __init__(self, fhalf, del_f):
        self.whalf = 2*pi*fhalf
        self.del_w = 2*pi*del_f

    def getParams(self):
        return (self.whalf, self.del_w)

class bode(object):
    def __init__(self, controlObj, cavBoxObj):
        self.Ki, self.Kp, self.tau = controlObj.getParams()
        self.whalf, self.del_w     = cavBoxObj.getParams()

    def setFrequency(self, f):
        self.f = f
        self.w = 2*pi*f
        self.s = 1j*self.w
    
    def propTransFunc(self):
        return self.Kp
    
    def intTransFunc(self):
        return self.Ki/self.s
    
    def controlTransFunc(self):
        return self.propTransFunc() + self.intTransFunc()
    
    def cavTransFunc(self):
        num   = self.whalf * (self.s + self.whalf + (1j*self.del_w))
        denom = self.del_w**2 + (self.s+self.whalf)**2
        return num/denom
    
    def delayTransFunc(self):
        return exp(-self.tau*self.s)
    
    def sysTransFunc(self):
        H   = self.cavTransFunc()
        E   = self.delayTransFunc()
        GH  = self.controlTransFunc() * H
        GHE = GH * E
        return H / (1 + GHE)
    
    def openLoop(self):
        H   = self.cavTransFunc()
        E   = self.delayTransFunc()
        GH  = self.controlTransFunc() * H
        return GH * E
    
f     = logspace(-1, 7, num=500)
tau   = 2e-6
fhalf = 518
Kp    = 50
Ki    = 1 * Kp * 2*pi*fhalf
del_f = 100

A = control(Kp=Kp, Ki=Ki, tau=tau)
B = cavBox(fhalf=fhalf, del_f=del_f)
bodeObj = bode(controlObj=A, cavBoxObj=B)
bodeObj.setFrequency(f)

fig = plt.figure(1)
ax211 = fig.add_subplot(211)
ax212 = fig.add_subplot(212)

for scale in logspace(-3, 3, 7):
    bodeObj.Ki = scale * Kp * 2*pi*fhalf
    resp = bodeObj.openLoop()
    ax211.semilogx(f, 20*log10(abs(resp)), label='%0.3fKi' % scale)
    ax212.semilogx(f, rad2deg(unwrap(map(phase, resp))), label='%0.3fKi' % scale)

ax211.set_ylabel('Noise suppression (dB)')
ax212.set_ylabel('Phase (rad)')

ax212.set_ylim(bottom=-360, top=10)

for ax in [ax211, ax212]:
    ax.set_xlabel('Frequency (Hz)')
    ax.grid()
    ax.legend()

plt.show()

